import argparse
import math
import os
import sys
import time

import pygame
from pygame.color import THECOLORS

parser = argparse.ArgumentParser()
parser.add_argument('-f', '--fps', help='plot peroid', type=int, default=24)
parser.add_argument('-s', '--stop', help='stop at end',
                    type=bool, default=True)
args = parser.parse_args()

pygame.init()
pygame.display.set_caption("RoboMaster")
os.environ['SDL_VIDEO_WINDOW_POS'] = "100,100"

size_fac = 30
FPS = args.fps
shape = (0, 0)
ashape = (0, 0)
bk = size_fac
rbk = 15*size_fac
robo_info = dict()
run_info = dict()
sup_info = dict()

inj2ammoSize = {'100': 0.15, '10': 0.07}
inj2color = {'100': 'mediumpurple2', '10': 'deeppink1'}
team2color = {'1': 'red', '2': 'blue'}
font_hight = int(0.5*size_fac)
font = pygame.font.SysFont('SimHei', font_hight)
showstart = 0
txt = []


def prase_loc(locstr):
    return tuple(int(float(x)*size_fac) for x in locstr.split(':'))


def preLoadPlt(fname):
    shape = (0, 0)
    robo_info = dict()
    run_info = dict()
    with open(fname, 'r') as f:
        lines = f.readlines()
        configs = list(filter(lambda x: x[:2] in ['bp', 'ar'], lines))
        runtime = list(filter(lambda x: x[:2] in [
                       'rb', 'am', 'ts', 'lu', 'au', 'rd', 'ks', 'ru', 'as', 'sp', 'rr'], lines))
        for line in configs:
            ctrl = line[:2].strip()
            info = line[3:].strip()
            infos = [x.strip() for x in info.split('@')]
            if ctrl == 'bp':
                shape = prase_loc(info)
            elif ctrl == 'ar':
                uid, team, name, maxlife, bps = infos
                bps = [prase_loc(x) for x in bps.split('|') if ':' in x]
                robo_info[uid] = (team, name, bps)
                run_info[uid] = [1., int(maxlife), '0', '0', '0', '0']
    return shape, robo_info, run_info, runtime


def drawLife(surf, loc, lif):
    pygame.draw.rect(surf, THECOLORS['green'],
                     (loc[0]-size_fac*0.5, loc[1]-size_fac*0.5, size_fac, size_fac*0.2))
    if lif < 1.0:
        pygame.draw.rect(surf, THECOLORS['red'],
                         (loc[0]-size_fac*0.5, loc[1]-size_fac*0.5, size_fac - lif * size_fac, size_fac*0.2))


def plotTxt(surf):
    surf.fill(THECOLORS['white'])
    txtToShow = font.render(
        ' ID    Name(Team) Life Num17 Heat17 Num42 Heat42', True, THECOLORS['black'])
    surf.blit(txtToShow, (0, 0))
    for i, (uid, robo) in enumerate(robo_info.items(), 1):
        team, name, _ = robo
        life, maxlife, a17num, a17heat, a42num, a42heat = run_info[uid]
        real_life = int(maxlife*life)
        real_life = str(real_life) if real_life < 10000 else 'Inf'
        txtToShow = font.render('{:>3s} {:>10s}({}) {:>4s} {:>5s} {:>6s} {:>5s} {:>6s}'.format(
            uid, name, team, real_life, a17num, a17heat, a42num, a42heat), True, THECOLORS['black'])
        surf.blit(txtToShow, (0, i*font_hight))

    for i, line in enumerate(reversed(txt[:-showstart] if showstart != 0 else txt), len(robo_info)+2):
        txtToShow = font.render(line, True, THECOLORS['black'])
        surf.blit(txtToShow, (0, i*font_hight))


def conv(gloc, grot, loc):
    cosOfAngle = math.cos(grot)
    sinOfAngle = math.sin(grot)
    return (gloc[0]+loc[0] * cosOfAngle - loc[1] * sinOfAngle,
            gloc[1]+loc[0] * sinOfAngle + loc[1] * cosOfAngle)


shape, robo_info, init_info, runtime = preLoadPlt('PlotDefault.plt')
run_info = init_info
ashape = (shape[0]+2*bk+rbk, shape[1]+2*bk)
screen = pygame.display.set_mode(ashape)
background = pygame.Surface(screen.get_size())
background = background.convert_alpha()
background.fill(THECOLORS['white'])
canvas = background.subsurface((bk, bk, shape[0], shape[1]))
canvas.fill(THECOLORS['black'])
alphaSur = pygame.Surface(shape)
alphaSur.set_alpha(128)
alphaSur.fill(THECOLORS['black'])
textArea = background.subsurface((shape[0]+2*bk, 0, rbk, shape[1]+2*bk))
runStatus = background.subsurface((0, 0, shape[0]+2*bk, bk))

runidx = 0
paused = False
curts = 0
clk = pygame.time.Clock()
lstSurf = canvas.copy()
while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT or event.type == pygame.KEYDOWN and (event.key == pygame.K_ESCAPE or event.key == pygame.K_q):
            pygame.quit()
            sys.exit(0)
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_DOWN or event.type == pygame.MOUSEBUTTONDOWN and event.button == 5:
            showstart += 5
            showstart = min(showstart, max(0, len(txt)-10))
            plotTxt(textArea)
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_UP or event.type == pygame.MOUSEBUTTONDOWN and event.button == 4:
            showstart -= 5
            showstart = max(showstart, 0)
            plotTxt(textArea)
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_LEFT:
            FPS /= 2.
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_RIGHT:
            FPS *= 2.
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_p:
            paused = not paused
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_r:
            curts = 0
            runidx = 0
            run_info = init_info

    if runidx == len(runtime) or paused:
        background.blit(lstSurf, (bk, bk))
        screen.blit(background, (0, 0))
        pygame.display.update()
        continue

    line = runtime[runidx]
    runidx += 1
    if runidx == len(runtime):
        txt.append('End.')
        plotTxt(textArea)

    ctrl = line[:2].strip()
    info = line[3:].strip()
    infos = [x.strip() for x in info.split('@')]

    if ctrl == 'ts':
        curts = int(int(info)/1e6+0.4999999)
        plotTxt(textArea)
        tp = clk.tick(FPS)
        txtToShow = font.render(
            ' Time: {} ms    FPS: {}'.format(int(int(info)/1e3+0.4999), int(1000./tp)), True, THECOLORS['black'])
        runStatus.fill(THECOLORS['white'])
        runStatus.blit(txtToShow, (0, 0))
        screen.blit(background, (0, 0))
        pygame.display.update()
        lstSurf = canvas.copy()
        canvas.fill(THECOLORS['black'])
        canvas.blit(alphaSur, (0, 0))
    elif ctrl == 'rb':
        uid, loc, life, level, a17num, a17heat, a42num, a42heat = infos
        loc, phi = loc.split('|')
        loc = prase_loc(loc)
        phi = float(phi)
        life = float(life)
        team, name, bps = robo_info[uid]
        bps = [conv(loc, phi, each) for each in bps]
        run_info[uid][0] = life
        run_info[uid][2:] = a17num, a17heat, a42num, a42heat
        if life > 0.01:
            drawLife(canvas, loc, life)
            pygame.draw.polygon(
                canvas, THECOLORS[team2color[team]], bps)
        else:
            pygame.draw.polygon(canvas, THECOLORS["gray"], bps)
    elif ctrl == 'am':
        inj, loc = infos
        loc = prase_loc(loc)
        pygame.draw.circle(
            canvas, THECOLORS[inj2color[inj]], loc, int(inj2ammoSize[inj]*size_fac))
    elif ctrl == 'lu':
        uid, level, cur = infos
        team, name, _ = robo_info[uid]
        run_info[uid][1] = int(cur)
        if 'UVA' not in name and 'Tower' not in name:
            txt.append('At {}: {}({}) upgrade max life level to {}({})'.format(
                curts, name, team, level, cur))
    elif ctrl == 'au':
        uid, level, _, _ = infos
        team, name, _ = robo_info[uid]
        if 'UVA' not in name and 'Tower' not in name:
            txt.append('At {}: {}({}) upgrade max ammo speed level to {}'.format(
                curts, name, team, level))
    elif ctrl == 'rd':
        target, shooter = infos
        shooter = shooter.split('|')
        team, name, _ = robo_info[target]
        txt.append('At {}: {}({}) destory by {}'.format(
            curts, name, team, shooter[0]) + ' and ' + ','.join(shooter[1:]))
    elif ctrl == 'ks':
        team, name, _ = robo_info[info]
        txt.append('At {}: {}({}) kill itself'.format(curts, name, team))
    elif ctrl == 'ru':
        uid, level = infos
        team, name, _ = robo_info[uid]
        if 'UVA' not in name and 'Tower' not in name:
            txt.append('At {}: {}({}) upgrade to {}'.format(
                curts, name, team, level))
    elif ctrl == 'as':
        team, inc, loc = infos
        bps = [prase_loc(x) for x in loc.split('|') if ':' in x]
        sup_info[team] = (inc, bps)
        pygame.draw.polygon(alphaSur, THECOLORS[team2color[team]], bps)
    elif ctrl == 'sp':
        team, n17, n42 = infos
        inc, bps = sup_info[team]
    elif ctrl == 'rr':
        src, target = infos
        team, name, _ = robo_info[target]
        if src == '255':
            txt.append('At {}: {}({}) resurrect itself'.format(
                curts, name, team))
        else:
            txt.append('At {}: {}({}) resurrect by engineer'.format(
                curts, name, team))
